
#include "mla_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <cmath>
#include <numeric>
#include <vector>
#include <algorithm>

constexpr int32_t TILING_HEAD_SIZE = 18;
constexpr int32_t TILING_PARA_SIZE = 8;
constexpr int32_t DIM_0 = 0;
constexpr int32_t DIM_1 = 1;
constexpr int32_t DIM_2 = 2;
constexpr int32_t DIM_3 = 3;
constexpr int32_t DIM_4 = 4;

namespace optiling {
const int32_t NUM0 = 0;
const int32_t NUM1 = 1;
const int32_t NUM2 = 2;
const int32_t NUM3 = 3;
const int32_t NUM4 = 4;
const int32_t NUM5 = 5;
const int32_t NUM6 = 6;
const int32_t NUM7 = 7;
const int32_t NUM8 = 8;
const int32_t NUM9 = 9;
const int32_t NUM10 = 10;
const int32_t NUM11 = 11;
const int32_t NUM12 = 12;
const int32_t NUM13 = 13;
const int32_t NUM14 = 14;
const int32_t NUM15 = 15;
const int32_t NUM16 = 16;
const int32_t NUM17 = 17;
const int32_t NUM32 = 32;
const int32_t NUM64 = 64;
const int32_t NUM128 = 128;
const int32_t NUM256 = 256;
const int32_t NUM512 = 512;
const int32_t NUM576 = 576;
const int32_t INDEX125 = 125;
const int32_t INDEX126 = 126;
const int32_t INDEX127 = 127;
const int32_t INDEX190 = 190;
const int32_t INDEX191 = 191;
const int32_t SPECIALNUM_TOKENS = 16;
const int32_t SPECIALNUM_HEADS = 32;
const int32_t EMBEDDING_LIMIT = 128;
const int32_t HIGH_32BIT = 32;
const uint32_t BATCH_MLA = 32;
const uint32_t BLOCK_DIM_MLA = 20;
const int32_t PP_MM_NUM = 8;
const int32_t PP_INDEX = 16;
constexpr int32_t M_LIMIT = 128;
const float SPLITKV_SEQLEN = 2048;
constexpr int32_t WORKSPACE_BLOCK_SIZE_DB = 65536;

constexpr std::array<int32_t, PP_MM_NUM> PP_MM = { 16, 32, 48, 64, 80, 96, 112, 128 };
constexpr std::array<int32_t, NUM6> QN_TILE_LIST = { 128, 64, 32, 16, 8, 1 };

using IndexArr = std::array<int32_t, NUM4>;

inline uint32_t GetHigh32Bit(uint64_t v) { return static_cast<uint32_t>(v >> HIGH_32BIT); }
inline uint32_t GetLoww32Bit(uint64_t v) { return static_cast<uint32_t>(v); }

enum CalcType : int {
    CALC_TYPE_UNDEFINED = 0, // 默认值
    CALC_TYPE_SPEC,          // 支持传入大于1的qseqlen
    CALC_TYPE_RING,          // ringAttention
    CALC_TYPE_SPEC_AND_RING,          // 支持传入大于1的qseqlen ringAttention
    CALC_TYPE_PREFILL,       // 全量场景
};

enum TilingKeyType : int{
    TILING_HALF_DATA = 0,
    TILING_BF16_DATA = 1,
    TILING_INT8_HALF_DATA = 2,
    TILING_INT8_BF16_DATA = 3
};

using AddrOffsets = struct AddressOffsetInfo {
    uint64_t addrQSeqOffset = 0;
    uint64_t addrKSeqOffset = 0;
    uint64_t addrVSeqOffset = 0;
    uint64_t addrOSeqOffset = 0;
    uint64_t addrOFdSeqOffset = 0;
    uint64_t addrLSeqOffset = 0;
    uint64_t addrMaskOffset = 0;
    int32_t totalQBlkNum = 0;
    int32_t block = 0;
};

// void GetAddrOffsetMLA(uint32_t *tilingParam, const AddrOffsets addrOffsets, const int32_t tilingOffset)
// {
//     tilingParam[tilingOffset + NUM2] = GetHigh32Bit(addrOffsets.addrQSeqOffset);
//     tilingParam[tilingOffset + NUM3] = GetLoww32Bit(addrOffsets.addrQSeqOffset);
//     tilingParam[tilingOffset + NUM4] = GetHigh32Bit(addrOffsets.addrOSeqOffset);
//     tilingParam[tilingOffset + NUM5] = GetLoww32Bit(addrOffsets.addrOSeqOffset);

//     // mask offset
//     tilingParam[tilingOffset + NUM6] = GetHigh32Bit(addrOffsets.addrMaskOffset);
//     tilingParam[tilingOffset + NUM7] = GetLoww32Bit(addrOffsets.addrMaskOffset);
// }

int32_t CalcSplitNum(uint32_t flashDecodingTaskNum, bool quantFlag, int32_t blockDim, int32_t minKVSeqlen, int32_t blockSize)
{
    if (blockDim - flashDecodingTaskNum <= NUM4 || quantFlag) {
        return NUM1;
    }
    if (blockSize == 0 || blockDim == 0) {
        return NUM1;
    }
    int32_t minKVBlocks = (minKVSeqlen + blockSize - 1) / blockSize;
    for (int32_t splitNum = 2; splitNum <= NUM6; splitNum++) {
        if ((flashDecodingTaskNum * splitNum) % blockDim != 0) {
            continue;
        }
        int32_t repeatTimesPerBlock = flashDecodingTaskNum * splitNum / blockDim;
        if (minKVSeqlen / splitNum >= blockSize &&
            repeatTimesPerBlock + NUM6 <= minKVBlocks - (minKVBlocks + splitNum - 1) / splitNum * repeatTimesPerBlock) {
            return splitNum;
        } else {
            return NUM1;
        }
    }
    return NUM1;
}
int32_t GetQNBlockTile(uint32_t numHeads, uint32_t kvHeads, int32_t qSeqLen)
{
    int32_t tileListIdx = static_cast<int32_t>(std::ceil(std::log2(qSeqLen)));
    tileListIdx = (tileListIdx > NUM5) ? NUM5 : tileListIdx;
    int32_t qNBlockTile = QN_TILE_LIST[tileListIdx];
    int32_t group = numHeads / kvHeads;
    qNBlockTile = (qNBlockTile > group) ? group : qNBlockTile;

    return qNBlockTile;
}

static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    MLATilingData tiling;

    auto *pkvSeqLencv = context->GetAttrs()->GetListInt(8);
    const int64_t *pkvSeqLen = pkvSeqLencv->GetData();
    int32_t kvSeqLenSize = pkvSeqLencv->GetSize();
    uint32_t batch = kvSeqLenSize;
    std::vector<uint32_t> kvSeqLenVec;
    for(int32_t k = 0; k < batch; k++) {
        kvSeqLenVec.push_back(*(pkvSeqLen + k));
    }
    auto maxKvSeq = *std::max_element(kvSeqLenVec.begin(), kvSeqLenVec.end()); 
    uint32_t maxKVSeqLen = maxKvSeq;
    tiling.set_num_batches(kvSeqLenSize);
    tiling.set_maxKVSeqLen(maxKVSeqLen);

    tiling.set_tiling_head_size(TILING_HEAD_SIZE);
    tiling.set_tiling_para_size(TILING_PARA_SIZE);

    auto *pheadSize = context->GetAttrs()->GetInt(1); 
    uint32_t q_heads = *pheadSize;
    uint32_t numHeads = *pheadSize;
    tiling.set_q_heads(q_heads);

    const gert::StorageShape* kcacheShape = context->GetInputShape(DIM_2);
    auto KDims = kcacheShape->GetStorageShape().GetDimNum();
    bool kNz = (kcacheShape->GetStorageShape().GetDim(KDims - 1)  == NUM16 || kcacheShape->GetStorageShape().GetDim(KDims - 1) == NUM32) ? 1 : 0;
    uint32_t embedding_size = 0;
    uint32_t block_size = 0;
    if (kNz) {
        embedding_size = static_cast<int32_t>(kcacheShape->GetStorageShape().GetDim(DIM_3)) *
                            static_cast<int32_t>(kcacheShape->GetStorageShape().GetDim(DIM_1));
        block_size= static_cast<int32_t>(kcacheShape->GetStorageShape().GetDim(DIM_2));
    } else {
        embedding_size= static_cast<int32_t>(kcacheShape->GetStorageShape().GetDim(DIM_3));
        block_size = static_cast<int32_t>(kcacheShape->GetStorageShape().GetDim(DIM_1));
    }
    tiling.set_embedding_size(embedding_size);
    tiling.set_block_size(block_size);

    uint32_t max_num_blocks_per_query = static_cast<int32_t>(context->GetInputShape(DIM_4)->GetStorageShape().GetDim(DIM_1));
    tiling.set_max_num_blocks_per_query(max_num_blocks_per_query);

    auto *ptor = context->GetAttrs()->GetFloat(2);
    float tor = *ptor;
    tiling.set_tor(tor);

    auto *pkvHeads = context->GetAttrs()->GetInt(3);
    int32_t kv_heads = *pkvHeads <= 0 ? *pheadSize : *pkvHeads;
    tiling.set_kv_heads(kv_heads);

    int32_t maxQSeqlen = 1;
    int32_t minQSeqlen = 1;

    auto *pqSeqLencv = context->GetAttrs()->GetListInt(7);
    const int64_t *pqSeqLen = pqSeqLencv->GetData();
    int32_t qSeqLenSize = pqSeqLencv->GetSize();
    int32_t minQSeqlentmp = 2147483647;
    int32_t maxQSeqlentmp = -2147483648;
    std::vector<int32_t> qSeqLenVec;
    for(int32_t k = 0; k < qSeqLenSize; k++) {
        minQSeqlentmp = minQSeqlentmp < *(pqSeqLen + k) ? minQSeqlentmp :  *(pqSeqLen + k);
        maxQSeqlentmp = maxQSeqlentmp > *(pqSeqLen + k) ? maxQSeqlentmp :  *(pqSeqLen + k);
        qSeqLenVec.push_back(*(pqSeqLen + k));
    }
    maxQSeqlen = maxQSeqlentmp;
    minQSeqlen = minQSeqlentmp;
    
    int32_t minKVSeqlen = 2147483647;
    for(int32_t k = 0; k < kvSeqLenSize; k++) {
        minKVSeqlen = minKVSeqlen < *(pkvSeqLen + k) ? minKVSeqlen : *(pkvSeqLen + k);
    }

    int32_t curQNBlockTile = GetQNBlockTile(q_heads, kv_heads, maxQSeqlen);
    tiling.set_cur_qn_blk_size(curQNBlockTile);

    auto *pmaskType = context->GetAttrs()->GetInt(5); 
    uint32_t maskType = *pmaskType;
    tiling.set_mask_type(maskType);

    auto *ptype = context->GetAttrs()->GetInt(0);
    bool quantFlag = (*ptype < NUM2) ? false : true;
    bool mtpTp1Flag = (*pheadSize == M_LIMIT);
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint32_t blockDim = ascendcPlatform.GetCoreNumAic();

    
    // uint32_t totalTaskNum = qSeqLenVec != nullptr ?
    //                       std::accumulate(qSeqLenVec.begin(), qSeqLenVec.end(), static_cast<int32_t>(0)) :
    //                       batch;
    uint32_t totalTaskNum = batch;

    uint32_t tailTaskNum = totalTaskNum % blockDim;
    uint32_t tailBatch = batch % blockDim;
    uint32_t flashDecodingTaskNum = quantFlag ? tailTaskNum : tailBatch;

    const int64_t *pisRing = context->GetAttrs()->GetInt(9);
    int32_t isRing = *pisRing;

    uint32_t dataType = 0;
    if (context->GetInputDesc(DIM_0)->GetDataType() == ge::DT_BF16) {
        dataType = TilingKeyType::TILING_BF16_DATA;
    } else if (context->GetInputDesc(DIM_0)->GetDataType() == ge::DT_FLOAT16) {
        dataType = TilingKeyType::TILING_HALF_DATA;
    } else if (context->GetInputDesc(DIM_1)->GetDataType() == ge::DT_FLOAT16) {
        dataType = TilingKeyType::TILING_INT8_HALF_DATA;
    } else {    
        dataType = TilingKeyType::TILING_INT8_BF16_DATA;
    } 

    bool flashDecoding = !quantFlag && flashDecodingTaskNum != 0 && isRing == 0 &&
                        minKVSeqlen >= SPLITKV_SEQLEN &&
                        ((minQSeqlen == NUM2 && maxQSeqlen == NUM2) ||
                        (minQSeqlen == 1 && maxQSeqlen == 1));

    uint32_t tilingKey = dataType + (kNz << NUM4) + (mtpTp1Flag << NUM2) +
                         (isRing << NUM5) + (flashDecoding << NUM6);

    tiling.set_totalTaskNum(totalTaskNum);


    uint32_t splitKVNum = blockDim / flashDecodingTaskNum > 1 ?  blockDim / flashDecodingTaskNum :
                    CalcSplitNum(flashDecodingTaskNum, quantFlag, blockDim, minKVSeqlen, block_size);
    tiling.set_kv_split_core_num(splitKVNum);                

    AddrOffsets addrOffsets {};
    // uint32_t emptySeq = (qSeqLen == nullptr) ? 1 : 0;
    uint32_t emptySeq = 1;
    uint32_t q_seqlen[32];
    uint32_t kv_seqlen[32];
    uint32_t addr_q_high32[32];
    uint32_t addr_q_loww32[32];
    uint32_t addr_o_high32[32];
    uint32_t addr_o_loww32[32];
    uint32_t addr_mask_high32[32];
    uint32_t addr_mask_loww32[32];
    for (int32_t seqIdx = 0; seqIdx < batch; seqIdx++) {
        int32_t qSeqLen = 1;
        qSeqLen = (emptySeq == 1) ? 1 : qSeqLenVec[seqIdx];
        qSeqLen = (kvSeqLenVec[seqIdx] == 0) ? 0 : qSeqLen;
        int32_t kvSeqlen = kvSeqLenVec[seqIdx];
        std::cout << "kvSeqlen:" << kvSeqlen << std::endl;

        // int32_t tilingOffset = TILING_HEAD_SIZE + TILING_PARA_SIZE * seqIdx;
        // tilingParam[tilingOffset] = static_cast<uint32_t>(qSeqLen);
        // tilingParam[tilingOffset + 1] = static_cast<uint32_t>(kvSeqlen);
        q_seqlen[seqIdx] = qSeqLen;
        kv_seqlen[seqIdx] = kvSeqlen;

        // GetAddrOffsetMLA(tilingParam, addrOffsets, tilingOffset);
        uint64_t addressQffset = static_cast<uint64_t>(q_heads * qSeqLen);
        uint64_t addressOffset = static_cast<uint64_t>(q_heads * embedding_size * qSeqLen);
        uint64_t addressMaskOffset = static_cast<uint64_t>(qSeqLen * maxKVSeqLen);
        addrOffsets.addrQSeqOffset += addressQffset;
        addrOffsets.addrOSeqOffset += addressOffset;
        addrOffsets.addrMaskOffset += addressMaskOffset;
        addr_q_high32[seqIdx] = GetHigh32Bit(addrOffsets.addrQSeqOffset);
        addr_q_loww32[seqIdx] = GetLoww32Bit(addrOffsets.addrQSeqOffset);
        addr_o_high32[seqIdx] = GetHigh32Bit(addrOffsets.addrOSeqOffset);
        addr_o_loww32[seqIdx] = GetLoww32Bit(addrOffsets.addrOSeqOffset);

        // mask offset
        addr_mask_high32[seqIdx] = GetHigh32Bit(addrOffsets.addrMaskOffset);
        addr_mask_loww32[seqIdx] = GetLoww32Bit(addrOffsets.addrMaskOffset);

    } 
    tiling.set_q_seqlen(q_seqlen);   
    tiling.set_kv_seqlen(kv_seqlen); 
    tiling.set_addr_q_high32(addr_q_high32); 
    tiling.set_addr_q_loww32(addr_q_loww32); 
    tiling.set_addr_o_high32(addr_o_high32); 
    tiling.set_addr_o_loww32(addr_o_loww32); 
    tiling.set_addr_mask_high32(addr_mask_high32); 
    tiling.set_addr_mask_loww32(addr_mask_loww32); 

    uint32_t dataLenHalf = sizeof(uint16_t);
    uint32_t dataLenFloat = sizeof(float);
    uint32_t dataLenInt = sizeof(int32_t);
    uint64_t basicWorkSpaceHalf = blockDim * WORKSPACE_BLOCK_SIZE_DB * dataLenHalf;
    uint64_t basicWorkSpaceFloat = blockDim * WORKSPACE_BLOCK_SIZE_DB * dataLenFloat;
    uint64_t basicWorkSpaceInt8 = blockDim * WORKSPACE_BLOCK_SIZE_DB * dataLenInt;
    uint64_t oCoreWorkSpaceSize = flashDecoding && mtpTp1Flag ?
        WORKSPACE_BLOCK_SIZE_DB * flashDecodingTaskNum * splitKVNum * dataLenFloat * 2 : 0;
    uint64_t lWorkSpaceSize = flashDecoding && mtpTp1Flag ?
        numHeads * flashDecodingTaskNum * splitKVNum * dataLenFloat * 2 * 8 : 0;
    size_t usrSize = 0;
    if (quantFlag) {
        uint64_t sWorkSpaceSize = mtpTp1Flag ? basicWorkSpaceFloat * 2 : basicWorkSpaceFloat;
        uint64_t pWorkSpaceSize = basicWorkSpaceInt8;
        uint64_t oTempWorkSpaceSize = basicWorkSpaceInt8 * 2;
        // kernelInfo.GetScratchSizes() = {sWorkSpaceSize, sWorkSpaceSize, pWorkSpaceSize,
        //                                 oTempWorkSpaceSize, basicWorkSpaceFloat, oCoreWorkSpaceSize, lWorkSpaceSize};
        usrSize = sWorkSpaceSize + sWorkSpaceSize + pWorkSpaceSize +       
                  oTempWorkSpaceSize + basicWorkSpaceFloat + oCoreWorkSpaceSize + lWorkSpaceSize;    
        tiling.set_sWorkSpaceSize(sWorkSpaceSize); 
        tiling.set_pWorkSpaceSize(pWorkSpaceSize); 
        tiling.set_oTempWorkSpaceSize(oTempWorkSpaceSize); 
        tiling.set_basicWorkSpaceFloat(basicWorkSpaceFloat); 
        tiling.set_oCoreWorkSpaceSize(oCoreWorkSpaceSize); 
        tiling.set_lWorkSpaceSize(lWorkSpaceSize);                                
    } else {
        uint64_t sWorkSpaceSize = mtpTp1Flag ? basicWorkSpaceFloat * 4 : basicWorkSpaceFloat * 2;
        uint64_t pWorkSpaceSize = mtpTp1Flag ? basicWorkSpaceHalf * 4 : basicWorkSpaceHalf * 2;
        uint64_t oTempWorkSpaceSize = mtpTp1Flag ? basicWorkSpaceFloat * 4 : basicWorkSpaceFloat * 2;
        uint64_t goWorkSpaceSize = mtpTp1Flag ? basicWorkSpaceFloat * 2 : basicWorkSpaceFloat;
        // kernelInfo.GetScratchSizes() = {sWorkSpaceSize, NUM512, pWorkSpaceSize, oTempWorkSpaceSize,
        //                                 goWorkSpaceSize, oCoreWorkSpaceSize, lWorkSpaceSize};
        usrSize = sWorkSpaceSize + NUM512 + pWorkSpaceSize + oTempWorkSpaceSize +
                  goWorkSpaceSize + oCoreWorkSpaceSize + lWorkSpaceSize;   
        tiling.set_sWorkSpaceSize(sWorkSpaceSize);
        tiling.set_pWorkSpaceSize(pWorkSpaceSize);
        tiling.set_oTempWorkSpaceSize(oTempWorkSpaceSize);
        tiling.set_goWorkSpaceSize(goWorkSpaceSize);
        tiling.set_oCoreWorkSpaceSize(oCoreWorkSpaceSize);
        tiling.set_lWorkSpaceSize(lWorkSpaceSize);                                       
    }
    std::cout << "usrSize:" << usrSize << std::endl;
    uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = usrSize + sysWorkspaceSize;


    context->SetTilingKey(tilingKey);
    context->SetBlockDim(blockDim);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    std::cout << "sWorkSpaceSize:" << tiling.get_sWorkSpaceSize() << std::endl;
    std::cout << "pWorkSpaceSize:" << tiling.get_pWorkSpaceSize() << std::endl;
    std::cout << "oTempWorkSpaceSize:" << tiling.get_oTempWorkSpaceSize() << std::endl;
    std::cout << "basicWorkSpaceFloat:" << tiling.get_basicWorkSpaceFloat() << std::endl;
    std::cout << "oCoreWorkSpaceSize:" << tiling.get_oCoreWorkSpaceSize() << std::endl;
    std::cout << "lWorkSpaceSize:" << tiling.get_lWorkSpaceSize() << std::endl;
    std::cout << "goWorkSpaceSize:" << tiling.get_goWorkSpaceSize() << std::endl;
    std::cout << "sSize:" << tiling.get_sSize() << std::endl;
    std::cout << "num_batches:" << tiling.get_num_batches() << std::endl;
    std::cout << "q_heads:" << tiling.get_q_heads() << std::endl;
    std::cout << "embedding_size:" << tiling.get_embedding_size() << std::endl;
    std::cout << "block_size:" << tiling.get_block_size() << std::endl;
    std::cout << "max_num_blocks_per_query:" << tiling.get_max_num_blocks_per_query() << std::endl;
    std::cout << "tor:" << tiling.get_tor() << std::endl;
    std::cout << "kv_heads:" << tiling.get_kv_heads() << std::endl;
    std::cout << "tiling_head_size:" << tiling.get_tiling_head_size() << std::endl;
    std::cout << "tiling_para_size:" << tiling.get_tiling_para_size() << std::endl;
    std::cout << "cur_qn_blk_size:" << tiling.get_cur_qn_blk_size() << std::endl;
    std::cout << "total_block_num:" << tiling.get_total_block_num() << std::endl;
    std::cout << "mask_type:" << tiling.get_mask_type() << std::endl;
    std::cout << "totalTaskNum:" << tiling.get_totalTaskNum() << std::endl;
    std::cout << "maxKVSeqLen:" << tiling.get_maxKVSeqLen() << std::endl;
    std::cout << "kv_split_core_num:" << tiling.get_kv_split_core_num() << std::endl;
    std::cout << "flashDecodingTaskNum:" << tiling.get_flashDecodingTaskNum() << std::endl;
    std::cout << "tailbatch:" << tiling.get_tailbatch() << std::endl;
    // std::cout << "q_seqlen[0]" << tiling.get_q_seqlen() << std::endl;


    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
const auto inputDataType = context->GetInputDataType(0);
context->SetOutputDataType(0, inputDataType);
return ge::GRAPH_SUCCESS;
}
}


namespace ops {
class MLA : public OpDef {
public:
    explicit MLA(const char* name) : OpDef(name)
    {
        this->Input("query")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("queryRope")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("kvCache")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ});
        this->Input("kvCacheRope")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ});
        this->Input("block_tables")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("contextLens")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("mask")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("qSeqlen")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("qkDescale")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("pvDescale")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("attenOut")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("lseOut")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("type").Int();
        this->Attr("headSize").Int();
        this->Attr("tor").Float();
        this->Attr("kvHead").Int();
        this->Attr("isTriuMask").AttrType(OPTIONAL).Int(0);
        this->Attr("maskType").AttrType(OPTIONAL).Int(0);
        this->Attr("quantType").AttrType(OPTIONAL).Int(0);
        this->Attr("qSeqLen").AttrType(OPTIONAL).ListInt();
        this->Attr("kvSeqLen").AttrType(OPTIONAL).ListInt();
        this->Attr("isRing").AttrType(OPTIONAL).Int(0);
        this->Attr("cacheMode").AttrType(OPTIONAL).Int(1);

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310p");

    }
};

OP_ADD(MLA);
}
